#!/usr/bin/env python3
"""
测试数据集分割脚本的GT文件后缀名处理功能
"""

import os
import tempfile
import shutil
from pathlib import Path

def test_gt_suffix_handling():
    """测试GT文件后缀名处理"""
    print("=== 测试GT文件后缀名处理 ===")

    # 创建临时目录
    with tempfile.TemporaryDirectory() as temp_dir:
        temp_path = Path(temp_dir)

        # 创建源目录结构
        source_dir = temp_path / "source"
        source_train_dir = source_dir / "train"
        source_images_dir = source_train_dir / "images"
        source_gt_dir = source_train_dir / "gt"

        source_images_dir.mkdir(parents=True)
        source_gt_dir.mkdir(parents=True)

        # 创建测试文件
        test_files = [
            ("image1.jpg", "image1.png"),  # 图像jpg, GT png
            ("image2.png", "image2.png"),  # 图像png, GT png
            ("image3.jpg", "image3.jpg"),  # 图像jpg, GT jpg
        ]

        # 创建pair列表文件
        pair_lines = []
        for img_name, gt_name in test_files:
            # 创建实际文件
            (source_images_dir / img_name).write_text("fake image")
            (source_gt_dir / gt_name).write_text("fake gt")

            # 添加到pair列表
            pair_lines.append(f"{source_images_dir / img_name} {source_gt_dir / gt_name}")

        # 写入pair列表文件
        pair_file = source_dir / "train_pair.lst"
        with open(pair_file, 'w') as f:
            f.write('\n'.join(pair_lines))

        # 设置输出目录
        output_dir = temp_path / "output"

        # 导入并运行分割函数
        import sys
        sys.path.insert(0, '/home/dd/working/WorkpieceDetection/examples/edge_detection_robustness/external_projects/DexiNed')

        from split_dataset import split_dataset_9_1

        # 运行分割
        train_list, val_list = split_dataset_9_1(
            str(source_dir),
            str(output_dir),
            train_ratio=0.7,
            random_seed=42
        )

        print(f"训练集列表: {train_list}")
        print(f"验证集列表: {val_list}")

        # 检查结果
        with open(train_list, 'r') as f:
            train_lines = f.readlines()

        with open(val_list, 'r') as f:
            val_lines = f.readlines()

        print(f"训练集包含 {len(train_lines)} 对")
        print(f"验证集包含 {len(val_lines)} 对")

        # 验证文件对是否正确
        all_lines = train_lines + val_lines
        for line in all_lines:
            parts = line.strip().split()
            if len(parts) == 2:
                img_path, gt_path = parts
                img_name = os.path.basename(img_path)
                gt_name = os.path.basename(gt_path)

                print(f"图像: {img_name}, GT: {gt_name}")

                # 验证文件存在
                assert os.path.exists(img_path), f"图像文件不存在: {img_path}"
                assert os.path.exists(gt_path), f"GT文件不存在: {gt_path}"

        print("✓ 所有文件对都正确处理!")
        print("✓ GT文件后缀名处理功能正常!")

if __name__ == "__main__":
    test_gt_suffix_handling()